#include <iostream>
#include <vector>
#include <string>
#include <set>
#include <Vector.hpp>
using namespace std;

class Solution
{
private:
    struct TrieNode
    {
        vector<TrieNode *> children;
        bool isEnd;
        TrieNode()
        {
            isEnd = false;
            children.resize(26, nullptr);
        }
    };

    vector<vector<int>> offset = {{0, 1}, {0, -1}, {1, 0}, {-1, 0}};

public:
    vector<string> findWords(vector<vector<char>> &board, vector<string> &words)
    {
        // creat tree for words
        TrieNode root;
        for (const string &word : words)
        {
            TrieNode *cur = &root;
            for (const char &c : word)
            {
                if (cur->children[c - 'a'] == nullptr)
                    cur->children[c - 'a'] = new TrieNode();
                cur = cur->children[c - 'a'];
            }
            cur->isEnd = true;
        }
        // search board
        set<string> res;
        for (int x = 0; x < board.size(); x++)
        {
            for (int y = 0; y < board[0].size(); y++)
            {
                // search from [x,y]
                vector<vector<bool>> visited(board.size(), vector<bool>(board[0].size(), false));
                dfs(board, x, y, &root, res, "", visited);
            }
        }
        return vector<string>(res.begin(), res.end());
    }

    void dfs(vector<vector<char>> &board, int x, int y, TrieNode *cur, set<string> &res, string word, vector<vector<bool>> &visited)
    {
        if (visited[x][y])
            return;
        char c = board[x][y];
        TrieNode *next = cur->children[c - 'a'];
        if (next == nullptr)
            return;
        word.push_back(c);
        visited[x][y] = true;
        if (next->isEnd)
        {
            res.insert(word);
            // can remove word from tree
        }
        for (int i = 0; i < 4; i++)
        {
            int newX = x + offset[i][0];
            int newY = y + offset[i][1];
            if (newX < 0 || newX >= board.size() || newY < 0 || newY >= board[0].size())
                continue;
            dfs(board, newX, newY, next, res, word, visited);
        }
        word.pop_back();
        visited[x][y] = false;
    }
};

int main()
{
    cout << "Code212" << endl;
    Solution s;
    vector<vector<char>> board = {{'o', 'a', 'a', 'n'}, {'e', 't', 'a', 'e'}, {'i', 'h', 'k', 'r'}, {'i', 'f', 'l', 'v'}};
    vector<string> words = {"oath", "pea", "eat", "rain"};
    // vector<vector<char>> board = {{'a'}};
    // vector<string> words = {"a"};
    vector<string> res = s.findWords(board, words);
    PrintVector(res);
    return 0;
}